from neo4j import GraphDatabase
from neo4j.graph import Path  # 适配 neo4j>=4.0
from typing import Dict, List, Any, Optional
from loguru import logger
import neo4j  # 用于版本检测

class Neo4jClient:
    """
    Neo4j数据库客户端，负责执行Cypher查询并提取子图
    """
    def __init__(self, uri: str, user: str, password: str, database: str = "neo4j"):
        self.uri = uri
        self.user = user
        self.password = password
        self.database = database
        self.driver = None
        self.connect()
    
    def connect(self):
        """
        连接到Neo4j数据库
        """
        try:
            self.driver = GraphDatabase.driver(self.uri, auth=(self.user, self.password))
            logger.info(f"已连接到Neo4j数据库: {self.uri}")
        except Exception as e:
            logger.error(f"连接Neo4j数据库失败: {str(e)}")
            raise e
    
    def test_connection(self) -> bool:
        """
        测试数据库连接是否正常
        """
        try:
            with self.driver.session(database=self.database) as session:
                result = session.run("RETURN 1 as test")
                return result.single()["test"] == 1
        except Exception as e:
            logger.error(f"测试Neo4j连接失败: {str(e)}")
            raise e
    
    def close(self):
        """
        关闭数据库连接
        """
        if self.driver:
            self.driver.close()
            logger.info("Neo4j连接已关闭")
    
    def execute_query(self, query: str, params: Dict = None) -> List[Dict]:
        """
        执行Cypher查询并返回结果
        
        Args:
            query: Cypher查询语句
            params: 查询参数
            
        Returns:
            查询结果列表
        """
        if params is None:
            params = {}
            
        try:
            with self.driver.session(database=self.database) as session:
                result = session.run(query, params)
                # return [record.data() for record in result] 
                ## 返回结构类似于：
                ## print(result)
                ## [{'p': [{'name': '索鞍锈蚀'}, 'may_occur', {'name': '索鞍'}, 'has_member', {'name': '悬索桥'}]}]
                # 修改后的命令为
                return result.graph()  # 返回完整的图结构数据
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}\n查询: {query}\n参数: {params}")
            raise e
    
    def extract_subgraph(self, query: str, params: Dict = None, limit: int = 100) -> Dict[str, Any]:
        """
        执行查询并提取子图数据，格式化为前端可视化所需的结构
        
        Args:
            query: Cypher查询语句
            params: 查询参数
            limit: 返回的最大节点数量
            
        Returns:
            包含节点和边的子图数据
        """
        if "LIMIT" not in query.upper():
            query += f" LIMIT {limit}"
            
        try:
            result = self.execute_query(query, params)
            
            nodes_set = set()
            edges_set = set()
            nodes = []
            edges = []
            
            for record in result:
                # 1. 处理 a, r, b 结构（常见于 MATCH (a)-[r]->(b) RETURN a, r, b）
                node_a = record.get('a')
                relation = record.get('r')
                node_b = record.get('b')
                
                if node_a is not None:
                    node_a_id = self._get_node_id(node_a)
                    if node_a_id and node_a_id not in nodes_set:
                        nodes_set.add(node_a_id)
                        nodes.append(self._serialize_node(node_a))
                
                if node_b is not None:
                    node_b_id = self._get_node_id(node_b)
                    if node_b_id and node_b_id not in nodes_set:
                        nodes_set.add(node_b_id)
                        nodes.append(self._serialize_node(node_b))
                
                if relation is not None and node_a is not None and node_b is not None:
                    source_id = self._get_node_id(node_a)
                    target_id = self._get_node_id(node_b)
                    if source_id and target_id:
                        edge_type = relation.get('type', 'UNKNOWN')
                        edge_id = (source_id, edge_type, target_id)
                        if edge_id not in edges_set:
                            edges_set.add(edge_id)
                            edges.append({
                                'source': source_id,
                                'target': target_id,
                                'type': edge_type,
                                'properties': {k: v for k, v in relation.items() if k != 'type'}
                            })
                
                # 2. 处理路径或其他复杂结果
                for key, value in record.items():
                    if isinstance(value, Path):  # neo4j.graph.Path
                        self._process_path(value, nodes, edges, nodes_set, edges_set)
                    elif isinstance(value, dict):
                        if "nodes" in value and "relationships" in value:
                            self._process_dict_path(value, nodes, edges, nodes_set, edges_set)
                        elif "id" in value or "name" in value:
                            self._process_node(value, nodes, edges, nodes_set, edges_set)
                        elif "source" in value and "target" in value:
                            self._process_edge(value, nodes, edges, nodes_set, edges_set)
            
            return {
                "nodes": nodes,
                "edges": edges
            }
        except Exception as e:
            logger.error(f"提取子图失败: {str(e)}\n查询: {query}")
            return {"nodes": [], "edges": []}
    
    def _get_node_id(self, node) -> str:
        """
        获取节点的唯一标识，优先使用 'name' 属性，若无则使用其他属性或节点 ID。
        """
        if node is None:
            return None
        if isinstance(node, dict):
            if 'name' in node:
                return node['name']
            if 'id' in node:
                return node['id']
            for key in node:
                if key not in ('name', 'id', 'labels') and node[key]:
                    return str(node[key])
            return f"node_{id(node)}"
        else:  # neo4j.graph.Node
            if 'name' in node:
                return node['name']
            node_id = getattr(node, 'element_id', None)
            if node_id:
                return node_id
            for key in node:
                if key != 'name' and node[key]:
                    return str(node[key])
            return f"{':'.join(getattr(node, 'labels', []))}_{getattr(node, 'element_id', id(node))}"
    
    def _serialize_node(self, node) -> Dict[str, Any]:
        """
        将节点序列化为字典，保留所有属性和标签。
        """
        if node is None:
            return {}
        if isinstance(node, dict):
            return {
                'id': self._get_node_id(node),
                'labels': node.get('labels', []),
                'properties': {k: v for k, v in node.items() if k not in ('id', 'labels')}
            }
        else:  # neo4j.graph.Node
            return {
                'id': self._get_node_id(node),
                'labels': list(getattr(node, 'labels', [])),
                'properties': dict(node) if hasattr(node, '__getitem__') else {}
            }
    
    def _process_path(self, path: Path, nodes: List[Dict], edges: List[Dict], nodes_set: set, edges_set: set):
        """
        处理 Neo4j 路径对象，提取节点和边。
        """
        if not path:
            return
        for node in path.nodes:
            node_id = self._get_node_id(node)
            if node_id and node_id not in nodes_set:
                nodes_set.add(node_id)
                nodes.append(self._serialize_node(node))
        for rel in path.relationships:
            source_id = self._get_node_id(rel.start_node)
            target_id = self._get_node_id(rel.end_node)
            if source_id and target_id:
                edge_id = (source_id, rel.type, target_id)
                if edge_id not in edges_set:
                    edges_set.add(edge_id)
                    edges.append({
                        'source': source_id,
                        'target': target_id,
                        'type': rel.type,
                        'properties': dict(rel) if hasattr(rel, '__getitem__') else {}
                    })
    
    def _process_dict_path(self, path_data: Dict, nodes: List[Dict], edges: List[Dict], nodes_set: set, edges_set: set):
        """
        处理字典格式的路径数据（包含 nodes 和 relationships）。
        """
        for node in path_data.get("nodes", []):
            if isinstance(node, dict) and ("id" in node or "name" in node):
                node_id = self._get_node_id(node)
                if node_id and node_id not in nodes_set:
                    nodes_set.add(node_id)
                    nodes.append(self._serialize_node(node))
        
        for rel in path_data.get("relationships", []):
            if isinstance(rel, dict) and "source" in rel and "target" in rel:
                source_id = rel["source"]
                target_id = rel["target"]
                edge_type = rel.get("type", "UNKNOWN")
                edge_id = (source_id, edge_type, target_id)
                if edge_id not in edges_set:
                    edges_set.add(edge_id)
                    edges.append({
                        "source": source_id,
                        "target": target_id,
                        "type": edge_type,
                        "properties": {k: v for k, v in rel.items() if k not in ("source", "target", "type")}
                    })
    
    def _process_node(self, node_data: Dict, nodes: List[Dict], edges: List[Dict], nodes_set: set, edges_set: set):
        """
        处理单个节点数据（字典格式）。
        """
        node_id = self._get_node_id(node_data)
        if node_id and node_id not in nodes_set:
            nodes_set.add(node_id)
            nodes.append(self._serialize_node(node_data))
    
    def _process_edge(self, edge_data: Dict, nodes: List[Dict], edges: List[Dict], nodes_set: set, edges_set: set):
        """
        处理单个边数据（字典格式）。
        """
        source = edge_data.get("source")
        target = edge_data.get("target")
        edge_type = edge_data.get("type", "UNKNOWN")
        if source and target:
            edge_id = (source, edge_type, target)
            if edge_id not in edges_set:
                edges_set.add(edge_id)
                edges.append({
                    "source": source,
                    "target": target,
                    "type": edge_type,
                    "properties": {k: v for k, v in edge_data.items() if k not in ("source", "target", "type")}
                })
    
    def get_entity_info(self, entity_name: str) -> Optional[Dict]:
        """
        获取实体的详细信息
        
        Args:
            entity_name: 实体名称
            
        Returns:
            实体信息字典或None（如果未找到）
        """
        query = """
        MATCH (n)
        WHERE n.name = $name OR n.id = $name
        RETURN n
        LIMIT 1
        """
        
        result = self.execute_query(query, {"name": entity_name})
        if result and "n" in result[0]:
            return result[0]["n"]
        return None
    
    def get_relationship_between_entities(self, source_entity: str, target_entity: str, max_path_length: int = 3) -> List[Dict]:
        """
        查找两个实体之间的关系路径
        
        Args:
            source_entity: 源实体名称
            target_entity: 目标实体名称
            max_path_length: 最大路径长度
            
        Returns:
            路径列表
        """
        query = f"""
        MATCH (source), (target),
        p = shortestPath((source)-[*1..{max_path_length}]-(target))
        WHERE source.name = $source_name OR source.id = $source_name
        AND (target.name = $target_name OR target.id = $target_name)
        RETURN p
        LIMIT 10
        """
        
        return self.execute_query(query, {
            "source_name": source_entity,
            "target_name": target_entity
        })
    
    def search_entities(self, keyword: str, limit: int = 10) -> List[Dict]:
        """
        搜索包含关键词的实体
        
        Args:
            keyword: 搜索关键词
            limit: 返回结果数量限制
            
        Returns:
            匹配的实体列表
        """
        query = """
        MATCH (n)
        WHERE n.name CONTAINS $keyword OR n.description CONTAINS $keyword
        RETURN n
        LIMIT $limit
        """
        
        result = self.execute_query(query, {"keyword": keyword, "limit": limit})
        return [record["n"] for record in result] if result else []    
    # def search_entities(self, keyword: str, limit: int = 10) -L# ist[Dict]:
    #     """
    #     搜索包含关键词的实体
        
    #     Args:
    #     #   keyword: 搜索关键词
    #         limit: 返回结果数量限制
            
    #     Returns:
    #     #   匹配的实体列表
    #     """
    #     quer=#  """
    #     MATCH (n)
    #     WHERE na# me CONTAINS $keyword OR n.description N# TAINS $keyword
    #     RETURN n
    #     LIMIT $limit #        """
        
    #     result = self.exu# te_query(query, {"keyword": keyword, "limit": limit})
    #     return [record["n"] for record in result]

    def get_graph_data(self, query: str, params: Dict = None) -> Dict[str, List[Dict]]:
        """
        执行Cypher查询并返回适合前端解析的图结构数据
        
        Args:
            query: Cypher查询语句
            params: 查询参数
            
        Returns:
            包含节点和关系的字典，格式如下：
            {
                'nodes': [{'id': '622', 'labels': ['DiseaseName'], 'properties': {'name': '索鞍锈蚀'}}, ...],
                'relationships': [{'id': '687', 'type': 'may_occur', 'startNode': '573', 'endNode': '622', 'properties': {}}, ...]
            }
        """
        if params is None:
            params = {}

        try:
            with self.driver.session(database=self.database) as session:
                result = session.run(query, params)
                graph = result.graph()

                nodes = []
                for node in graph.nodes:
                    # 提取 element_id，确保处理类似 '4:32b37dbe-964a-495e-ab32-bcdf260e4534:622' 的格式
                    element_id = node.element_id.split(':')[-1]
                    nodes.append({
                        'id': element_id,
                        'labels': list(node.labels),
                        'properties': dict(node._properties)
                    })

                relationships = []
                for rel in graph.relationships:
                    relationships.append({
                        'id': rel.element_id.split(':')[-1],
                        'type': rel.type,
                        'startNode': rel.start_node.element_id.split(':')[-1],
                        'endNode': rel.end_node.element_id.split(':')[-1],
                        'properties': dict(rel._properties)
                    })
                # 处理后，relationships的表述，将使用edges进行替换
                return {'nodes': nodes, 'edges': relationships}
        except Exception as e:
            logger.error(f"执行查询失败: {str(e)}\n查询: {query}\n参数: {params}")
            raise e
